#include "lexer.h"

Lexer::Lexer(const std::string& source) : source(source), position(0), line(1), column(1) {
    // 初始化关键字映射表
    keywords = {
        {"int", TokenType::INT},
        {"void", TokenType::VOID},
        {"if", TokenType::IF},
        {"else", TokenType::ELSE},
        {"while", TokenType::WHILE},
        {"break", TokenType::BREAK},
        {"continue", TokenType::CONTINUE},
        {"return", TokenType::RETURN}
    };
}

// 执行词法分析并返回所有Token
vector<Token> Lexer::tokenize() {
    std::vector<Token> tokens;
    Token token = getNextToken();

    // 循环获取所有Token直到文件结束
    while (token.type != TokenType::END_OF_FILE) {
        tokens.push_back(token);
        token = getNextToken();
    }
    tokens.push_back(token); // 添加EOF token
    return tokens;
}

// 前进一个字符
void Lexer::advance() {
    if (position < source.length()) {
        if (source[position] == '\n') {
            line++;
            column = 1;
        }
        else column++;
        position++;
    }
}

// 查看当前字符（不前进位置）
char Lexer::peek() {
    if (position < source.length()) return source[position];
    //结束
    return '\0';
}

// 查看下一个字符（不前进位置）
char Lexer::peekNext() {
    if (position + 1 < source.length()) return source[position + 1];
    return '\0';
}

// 跳过空白字符（包括空格、制表符、换行符等）
void Lexer::skipWhitespace() {
    while (position < source.length() && isspace(peek())) advance();
}

// 跳过单行注释
void Lexer::skipLineComment() {
    // 跳过"//"后的所有字符直到行尾
    while (position < source.length() && peek() != '\n') advance();
    // 跳过换行符
    if (peek() == '\n') advance();
}

// 跳过多行注释
void Lexer::skipBlockComment() {
    // 已处于 '/' 位置，且下一个是 '*'
    advance(); // 跳过 '/' (现在位置在 '*')
    advance(); // 跳过 '*' (现在位置在注释内容开始)

    // 循环直到找到"*/"
    while (position < source.length()) {
        if (peek() == '*' && peekNext() == '/') {
            advance(); // 跳过 '*'
            advance(); // 跳过 '/'
            return;
        }
        advance();
    }

    // 如果到达这里，说明注释未正确关闭
    errors.push_back("Error: Unclosed block comment at line " +
        std::to_string(line) + ", column " + std::to_string(column));
}

// 处理注释
void Lexer::handleComments() {
    // 循环处理连续的注释
    while (true) {
        //先跳过空白部分
        skipWhitespace();

        if (peek() == '/' && peekNext() == '/') {
            // 单行注释
            advance(); // 跳过第一个 '/'
            advance(); // 跳过第二个 '/'
            skipLineComment();
        }
        else if (peek() == '/' && peekNext() == '*') {
            // 多行注释
            advance(); // 跳过第一个 '/'
            skipBlockComment();
        }
        else {
            // 没有注释了，退出循环
            break;
        }
    }
}

// 解析标识符或关键字
Token Lexer::parseIdentifierOrKeyword() {
    int startLine = line;
    int startCol = column;
    std::string lexeme;

    // 标识符规则：以字母或下划线开头，后跟字母、数字或下划线
    while (position < source.length() && (isalpha(peek()) || peek() == '_')) {
        lexeme += peek();
        advance();
    }

    // 检查是否是映射中的关键字
    auto it = keywords.find(lexeme);
    if (it != keywords.end()) {
        return Token(it->second, lexeme, startLine, startCol);
    }

    //没找到则是标识符
    return Token(TokenType::IDENTIFIER, lexeme, startLine, startCol);
}

// 解析数字（支持负号）
Token Lexer::parseNumber() {
    int startLine = line;
    int startCol = column;
    std::string lexeme;

    // 检查负号（如果存在）
    if (peek() == '-') {
        // 检查负号是否是数字的一部分（后面紧跟数字）
        if (isdigit(peekNext())) {
            lexeme += '-';
            advance();
        } else {
            // 如果不是数字，则返回单独的负号运算符
            advance();
            return Token(TokenType::SUB, "-", startLine, startCol);
        }
    }

    // 解析数字部分
    while (position < source.length() && isdigit(peek())) {
        lexeme += peek();
        advance();
    }

    return Token(TokenType::NUMBER, lexeme, startLine, startCol);
}

// 解析运算符
Token Lexer::parseOperator() {
    int startLine = line;
    int startCol = column;
    char current = peek();

    // 处理双字符运算符
    switch (current) {
    case '=':
        advance();
        if (peek() == '=') {
            advance();
            return Token(TokenType::EQ, "==", startLine, startCol);
        }
        return Token(TokenType::ASSIGN, "=", startLine, startCol);

    case '|':
        advance();
        if (peek() == '|') {
            advance();
            return Token(TokenType::OR, "||", startLine, startCol);
        }
        else {
            std::string unknown(1, current);
            errors.push_back("Error: Unknown operator '" + unknown + "' at line " +
                std::to_string(startLine) + ", column " + std::to_string(startCol));
            return Token(TokenType::UNKNOWN, unknown, startLine, startCol);
        }
        break;

    case '&':
        advance();
        if (peek() == '&') {
            advance();
            return Token(TokenType::AND, "&&", startLine, startCol);
        }
        else {
            std::string unknown(1, current);
            errors.push_back("Error: Unknown operator '" + unknown + "' at line " +
                std::to_string(startLine) + ", column " + std::to_string(startCol));
            return Token(TokenType::UNKNOWN, unknown, startLine, startCol);
        }
        break;

    case '<':
        advance();
        if (peek() == '=') {
            advance();
            return Token(TokenType::LE, "<=", startLine, startCol);
        }
        return Token(TokenType::LT, "<", startLine, startCol);

    case '>':
        advance();
        if (peek() == '=') {
            advance();
            return Token(TokenType::GE, ">=", startLine, startCol);
        }
        return Token(TokenType::GT, ">", startLine, startCol);

    case '!':
        advance();
        if (peek() == '=') {
            advance();
            return Token(TokenType::NE, "!=", startLine, startCol);
        }
        return Token(TokenType::NOT, "!", startLine, startCol);

    case '+': advance(); return Token(TokenType::ADD, "+", startLine, startCol);
    case '-': advance(); return Token(TokenType::SUB, "-", startLine, startCol);
    case '*': advance(); return Token(TokenType::MUL, "*", startLine, startCol);
    case '/': advance(); return Token(TokenType::DIV, "/", startLine, startCol);
    case '%': advance(); return Token(TokenType::MOD, "%", startLine, startCol);
    default:
        std::string unknown(1, current);
        errors.push_back("Error: Unknown operator '" + unknown + "' at line " +
            std::to_string(startLine) + ", column " + std::to_string(startCol));
        advance();
        return Token(TokenType::UNKNOWN, unknown, startLine, startCol);
    }
}

// 获取下一个Token
Token Lexer::getNextToken() {
    // 处理空白和注释
    handleComments();

    // 检查是否到达文件末尾
    if (position >= source.length()) {
        return Token(TokenType::END_OF_FILE, "", line, column);
    }

    char current = peek();

    // 根据当前字符类型分派到不同的解析器
    if (isalpha(current) || current == '_') {
        return parseIdentifierOrKeyword();
    }
    else if (isdigit(current) || (current == '-' && isdigit(peekNext()))) {
        return parseNumber();
    }
    else if (current == ',' || current == ';' || current == '(' ||
        current == ')' || current == '{' || current == '}') {
        return parseDelimiter();
    }
    else if (current == '=' || current == '|' || current == '&' || current == '>' ||
        current == '<' || current == '!' || current == '+' || current == '-' ||
        current == '*' || current == '/' || current == '%') {
        return parseOperator();
    }
    else {
        // 记录未知字符错误
        std::string charStr(1, current);
        errors.push_back("Error: Unknown character '" + charStr + "' at line " +
            std::to_string(line) + ", column " + std::to_string(column));
        advance(); // 跳过未知字符继续解析
        return Token(TokenType::UNKNOWN, charStr, line, column);
    }
}

// 解析分隔符
Token Lexer::parseDelimiter() {
    int startLine = line;
    int startCol = column;
    char current = peek();
    advance(); // 消耗当前字符进行下个token获取
    switch (current) {
    case ',': return Token(TokenType::COMMA, ",", startLine, startCol);
    case ';': return Token(TokenType::SEMICOLON, ";", startLine, startCol);
    case '(': return Token(TokenType::LPAREN, "(", startLine, startCol);
    case ')': return Token(TokenType::RPAREN, ")", startLine, startCol);
    case '{': return Token(TokenType::LBRACE, "{", startLine, startCol);
    case '}': return Token(TokenType::RBRACE, "}", startLine, startCol);
    default:
        std::string unknown(1, current);
        return Token(TokenType::UNKNOWN, unknown, startLine, startCol);
    }
}

void printTokens(const std::vector<Token>& tokens) {
    std::cout << std::left << std::setw(15) << "Token Type"
        << std::setw(20) << "Lexeme"
        << std::setw(10) << "Line"
        << std::setw(10) << "Column" << std::endl;
    std::cout << std::setfill('-') << std::setw(55) << "" << std::setfill(' ') << std::endl;

    for (const auto& token : tokens) {
        std::string typeStr;
        switch (token.type) {
        case TokenType::INT: typeStr = "INT"; break;
        case TokenType::VOID: typeStr = "VOID"; break;
        case TokenType::IF: typeStr = "IF"; break;
        case TokenType::ELSE: typeStr = "ELSE"; break;
        case TokenType::WHILE: typeStr = "WHILE"; break;
        case TokenType::BREAK: typeStr = "BREAK"; break;
        case TokenType::CONTINUE: typeStr = "CONTINUE"; break;
        case TokenType::RETURN: typeStr = "RETURN"; break;
        case TokenType::IDENTIFIER: typeStr = "IDENTIFIER"; break;
        case TokenType::NUMBER: typeStr = "NUMBER"; break;
        case TokenType::ASSIGN: typeStr = "ASSIGN"; break;
        case TokenType::OR: typeStr = "OR"; break;
        case TokenType::AND: typeStr = "AND"; break;
        case TokenType::LT: typeStr = "LT"; break;
        case TokenType::GT: typeStr = "GT"; break;
        case TokenType::LE: typeStr = "LE"; break;
        case TokenType::GE: typeStr = "GE"; break;
        case TokenType::EQ: typeStr = "EQ"; break;
        case TokenType::NE: typeStr = "NE"; break;
        case TokenType::ADD: typeStr = "ADD"; break;
        case TokenType::SUB: typeStr = "SUB"; break;
        case TokenType::MUL: typeStr = "MUL"; break;
        case TokenType::DIV: typeStr = "DIV"; break;
        case TokenType::MOD: typeStr = "MOD"; break;
        case TokenType::NOT: typeStr = "NOT"; break;
        case TokenType::COMMA: typeStr = "COMMA"; break;
        case TokenType::SEMICOLON: typeStr = "SEMICOLON"; break;
        case TokenType::LPAREN: typeStr = "LPAREN"; break;
        case TokenType::RPAREN: typeStr = "RPAREN"; break;
        case TokenType::LBRACE: typeStr = "LBRACE"; break;
        case TokenType::RBRACE: typeStr = "RBRACE"; break;
        case TokenType::END_OF_FILE: typeStr = "EOF"; break;
        case TokenType::UNKNOWN: typeStr = "UNKNOWN"; break;
        }
        std::cout << std::setw(15) << typeStr
            << std::setw(20) << token.lexeme
            << std::setw(10) << token.line
            << std::setw(10) << token.column << std::endl;
    }
}